from datetime import datetime
import backtrader as bt
import matplotlib.pyplot as plt
import akshare as ak
import pandas as pd

plt.rcParams["font.sans-serif"] = ["SimHei"]
plt.rcParams["axes.unicode_minus"] = False

# 获取股票后复权数据
stock_hfq_df = ak.stock_zh_a_hist(symbol="000001", adjust="hfq").iloc[:, :6]
# 处理字段命名
stock_hfq_df.columns = ['date', 'open', 'close', 'high', 'low', 'volume']
# 将date设为索引
stock_hfq_df.index = pd.to_datetime(stock_hfq_df['date'])

class MyStrategy(bt.Strategy):
    """
    主策略程序
    """
    params = (
        ("macd_fast", 12),  # MACD快线周期
        ("macd_slow", 26),  # MACD慢线周期
        ("macd_signal", 9),  # MACD信号线周期
    )

    def __init__(self):
        self.data_close = self.datas[0].close
        self.order = None
        self.buy_price = None
        self.buy_comm = None
        
        # 初始化MACD指标
        self.macd = bt.indicators.MACDHisto(self.datas[0],
                                           period_me1=self.params.macd_fast,
                                           period_me2=self.params.macd_slow,
                                           period_signal=self.params.macd_signal)

    def next(self):
        if self.order:  
            return
        
        # 检查MACD金叉和死叉
        macd_line = self.macd.lines.macd  # MACD线（快线-慢线）
        signal_line = self.macd.lines.signal  # 信号线
        
        if macd_line[0] > signal_line[0] and macd_line[-1] <= signal_line[-1]:  # MACD金叉，买入
            self.order = self.buy(size=100)
        elif macd_line[0] < signal_line[0] and macd_line[-1] >= signal_line[-1]:  # MACD死叉，卖出
            self.order = self.sell(size=100)

# 初始化回测系统
cerebro = bt.Cerebro()

# 设置回测时间范围
start_date = datetime(2010, 4, 3)
end_date = datetime(2024, 5, 22)

# 加载数据到Cerebro
data = bt.feeds.PandasData(dataname=stock_hfq_df, fromdate=start_date, todate=end_date)
cerebro.adddata(data)

# 加载策略到Cerebro
cerebro.addstrategy(MyStrategy)

# 设置初始资金和手续费
start_cash = 1000000
cerebro.broker.setcash(start_cash)
cerebro.broker.setcommission(commission=0.002)

# 运行回测
cerebro.run()

# 获取回测结果并打印
port_value = cerebro.broker.getvalue()
pnl = port_value - start_cash
print(f"初始资金: {start_cash}\n回测期间：{start_date.strftime('%Y%m%d')}:{end_date.strftime('%Y%m%d')}")
print(f"总资金: {round(port_value, 2)}")
print(f"净收益: {round(pnl, 2)}")

# 绘制图表
cerebro.plot(style='candlestick')